#include "refine/smooth2.h"
#include "refine/window_pass.h"
#include <limits>
#include <iostream>
#include "imglib/imageclass.h"
#include <imglib/bitmap.h>
#include <imglib/bmpmorph.h>
#include <glib/g3/vector.h>

using std::clog;
using std::endl;
using std::flush;

namespace refine {

namespace smooth2 {

const params default_params = {2, 2, 1, 100, 100, false};

typedef grey_mask::mask::elem_type grey_mask_value;
typedef unsigned uint;
typedef g3::basic_vector<uint>		vectoru;
typedef g3::basic_vector<float>		vectorf;

const uint uchar_max = std::numeric_limits<unsigned char>::max();
const uint uchar_avg = uchar_max / 2;

template <typename Color>
inline vectoru v(Color const & c)
	{ return vectoru(c.b, c.g, c.r); }

// reconstructs foreground and background in a point and from them and image color
// updates alpha in that point
template <typename Raster, typename Matte>
class smoother
{
public:
	smoother(Raster const & src_, 
		Matte const & matte_, grey_mask::mask & next_matte_,
		img::bitmap & modifiable_, float max_condition_number_)
		: src(src_), matte(matte_), next_matte(next_matte_), modifiable(modifiable_),
		  max_condition_number(max_condition_number_)
		{ clear(); }
	void clear() {
		remove_all();
		checks = achanges = 0;
	}
	size_extent dim() const
		{ return src.dim(); }
	uint check_count() const
		{ return checks; }
	uint change_count() const
		{ return achanges; }
	bool selected(size_point const & p) const
		{ return modifiable[p]; }
	void add(size_point const & p)
		{ add_sample(v(src[p]), matte[p]); }
	void remove(size_point const & p)
		{ remove_sample(v(src[p]), matte[p]); }
	void remove_all() {
		n = sum_alpha = sum_sqr_alpha = 0;
		sum_I = sum_alpha_I = vectoru::nill;
	}
	void process(size_point const & p, size_rect const &) {
		++checks;
		grey_mask_value g;
		if (!classify(v(src[p]), g) || matte[p] == g)
			modifiable[p] = false;
		else {
			++achanges;
			next_matte[p] = g;
		}
	}

private:
	/// guesses alpha for the given color
	bool classify(vectorf const & I, grey_mask_value & g) const 
	{
		//         [ a b ]
		// A^T A = [ b c ]
		float a = sum_sqr_alpha;
		float b = uchar_max * sum_alpha - sum_sqr_alpha;
		float c = uchar_max * uchar_max * n + sum_sqr_alpha - 2 * uchar_max * sum_alpha;

		float discriminant = sqr(a - c) + 4 * b * b;
		if (sqr(max_condition_number + 1) * discriminant > sqr((max_condition_number - 1) * (a + c)))
			return false;

		float rdet = uchar_max / (a * c - b * b);
		assert (rdet > 0);

		vectorf sum_I = float(uchar_max) * vectorf(this->sum_I);
		vectorf sum_alpha_I = vectorf(this->sum_alpha_I);

		vectorf F = rdet * ((c + b) * sum_alpha_I - b * sum_I);
		vectorf B = rdet * (-(a + b) * sum_alpha_I + a * sum_I);
		vectorf D = F - B;

		float den = sqr(D);
		if (den == 0)
			g = uchar_avg;
		else {
			float num = scalar_product(I - B, D);
			if (num <= 0)
				g = 0;
			else if (num >= den)
				g = uchar_max;
			else 
				g = uchar_max * num / den + 0.5f;
		}
		return true;
	}
private:
	void add_sample(vectoru const & I, uint alpha) {
		++n;
		sum_alpha += alpha;
		sum_sqr_alpha += alpha * alpha;
		sum_I += I;
		sum_alpha_I += alpha * I;
	}
	void remove_sample(vectoru const & I, uint alpha) {
		--n;
		sum_alpha -= alpha;
		sum_sqr_alpha -= alpha * alpha;
		sum_I -= I;
		sum_alpha_I -= alpha * I;
	}
private:
	Raster const & src;
    Matte const & matte;
	grey_mask::mask & next_matte;
	img::bitmap & modifiable;
	float max_condition_number;
private:
	uint n, sum_alpha, sum_sqr_alpha;
	vectoru sum_I, sum_alpha_I;
	uint checks, achanges;
};

template <typename Raster, typename Matte>
static void perform_core(
	const Raster & src, 
	Matte & matte, 
	const img::bitmap & strict_fore, 
	const img::bitmap & strict_back,
	img::bitmap & modifiable, ///< map of pixels that may be changed during the following pass
	const params & p)
{
	if (!p.silent)
		clog << "Volosok(R) matting..." << endl;

	size_extent dim = src.dim();
	assert (dim == matte.dim());
	assert (dim == modifiable.dim());
	int_rect frame(0, 0, dim.sx, dim.sy);

	grey_mask::mask next_matte = matte;

	for (uint pass = 0; pass < p.passes_num; ++pass) {
		if (!p.silent)
			clog << "pass " << pass+1 << '/' << p.passes_num << flush;

		img::combine(modifiable, strict_fore, img::bit_and_not());
		img::combine(modifiable, strict_back, img::bit_and_not());
		smoother<Raster, Matte> sm(src, matte, next_matte, modifiable, p.max_condition_number);
		window_pass(sm, p.whs);

		if (!p.silent)
			clog << ": checked = " << sm.check_count() << ", alpha changes = " << sm.change_count() << endl;

		if (sm.change_count() == 0)
			break;

		matte = next_matte;

		img::dilate8(modifiable, p.whs);
	}
}

template <typename Raster, typename Matte>
static void perform_from_bit_mask(
	const Raster & src, 
	const bit_mask::mask & m, 
	Matte & matte, 
	const params & p)
{
	size_extent dim = src.dim();
	assert (dim == m.dim());
	assert (dim == matte.dim());

	for (typename Matte::iterator i = matte.begin(); i != matte.end(); ++i)
		*i = m[i] ? uchar_max : 0;

	img::bitmap strict_fore, strict_back;
	img::erode4(strict_fore, m, p.fore_depth);
	img::dilate4(strict_back, m, p.back_depth);
	strict_back.inverse();

	img::bitmap modifiable = m; 
	{
		uint w = 1 + p.whs * (p.max_condition_number - 1) / (p.max_condition_number + 1);
		img::dilate4(modifiable, w);
		img::bitmap tmp = m;
		tmp.inverse();
		img::dilate4(tmp, w);
		img::combine(modifiable, tmp, img::bit_and());
	}

	perform_core(src, matte, strict_fore, strict_back, modifiable, p);
}

template <typename Raster, typename Matte>
static void perform_from_grey_mask(
	const Raster & src, 
	Matte & matte, 
	const params & p)
{
	size_extent dim = matte.dim();

	img::bitmap not_min(dim), not_max(dim), strict_fore(dim), strict_back(dim);
	for (typename Matte::/*const_*/iterator i = matte.begin(); i != matte.end(); ++i) {
		not_min[i] = *i != 0;
		not_max[i] = *i != uchar_max;
		strict_fore[i] = *i > uchar_avg;
	}

	strict_back = strict_fore;
	img::erode4(strict_fore, p.fore_depth);
	img::dilate4(strict_back, p.back_depth);
	strict_back.inverse();

	img::bitmap modifiable;
	{
		uint w = 1 + p.whs * (p.max_condition_number - 1) / (p.max_condition_number + 1);
		img::dilate4(modifiable, not_min, w);
		img::bitmap tmp;
		img::dilate4(tmp, not_max, w);
		img::combine(modifiable, tmp, img::bit_and());
	}

	perform_core(src, matte, strict_fore, strict_back, modifiable, p);
}

// exported functions

void perform(
	const byte_image & src, 
	const bit_mask::mask & m, 
	grey_mask::mask & matte, 
	const params & p)
{
	matte.resize(src.dim());
	perform_from_bit_mask(src, m, matte, p);
}

void perform(
	const byte_image & src, 
	grey_mask::mask & matte, 
	const params & p)
{
	matte.resize(src.dim());
	perform_from_grey_mask(src, matte, p);
}

void perform(
	byte_image & im, 
	const params & p)
{
	perform_from_grey_mask(im, img::projection(im, &byte_color::a), p);
}

void perform(
	argb_byte_image & im, 
	const params & p)
{
	perform_from_grey_mask(im, img::projection(im, &argb_byte_color::a), p);
}

} //namespace smooth2

} //namespace refine
